from skimage import io
import matplotlib.pyplot as plt
import os
from skimage.measure import block_reduce
import random
import networkx as nx
import matplotlib.pyplot as plt
import torch

from sensor import *
from graph_generator import *


from parameter import *
if not train_mode:
    from test_parameter import *

np.set_printoptions(threshold=np.inf)

class Env():
    def __init__(self, map_index, num_robots, random_seed=None, k_size=20, plot=False, test=False, input_type=INPUT_TYPE):
        # import environment ground truth from dungeon files
        self.test = test
        self.input_type = input_type
        if self.input_type == 'map':
            if self.test:
                self.map_dir = f'DungeonMaps/test_1025'  # change to 'complex', 'medium', and 'easy'
            else:
                self.map_dir = f'DungeonMaps/train_1025'
            self.random_seed = random_seed
            self.map_list = os.listdir(self.map_dir)
            self.map_list.sort(reverse=False)
            self.map_index = map_index % np.size(self.map_list)
            # self.ground_truth, self.start_position, self.target_position = self.import_ground_truth(
            #     self.map_dir + '/' + self.map_list[self.map_index])
            self.ground_truth = self.import_ground_truth(self.map_dir + '/' + self.map_list[self.map_index])

            self.ground_truth_size = np.shape(self.ground_truth)  # (480, 640)
            self.num_robots = num_robots
            # self.robot_positions = []
            # self.robot_indexes = []
            # for k in range(num_robots):
            #     self.robot_positions.append(self.start_position)
            # self.target_index = -1
            # initialize graph generator
            self.graph_generator = Graph_generator(map_size=self.ground_truth_size, k_size=k_size, plot=plot, test=self.test)
            if FIXED_OPPONENT == False:
                # self.node_coords, self.graph, self.node_feature, self.reference_policy, self.adjacent_matrix, self.network_adjacent_matrix, self.opponent_policy = None, None, None, None, None, None, None
                self.robot_positions, self.robot_indexes, self.target_position, self.target_index, self.node_coords, self.graph, self.node_feature, self.reference_policy, self.adjacent_matrix, self.network_adjacent_matrix, self.opponent_policy = None, None, None, None, None, None, None, None, None, None, None
            else:
                # self.node_coords, self.graph, self.node_feature, self.reference_policy, self.adjacent_matrix, self.network_adjacent_matrix = None, None, None, None, None, None
                self.robot_positions, self.robot_indexes, self.target_position, self.target_index, self.node_coords, self.graph, self.node_feature, self.reference_policy, self.adjacent_matrix, self.network_adjacent_matrix = None, None, None, None, None, None, None, None, None, None
            self.stepi = 0
            self.node_num = 0

            self.begin()
            self.robot_positions = copy.deepcopy(self.start_robot_positions)
            self.robot_indexes = copy.deepcopy(self.start_robot_indexes)
            self.target_position = copy.deepcopy(self.start_target_position)
            self.target_index = copy.deepcopy(self.start_target_index)

            k_size = self.node_coords.shape[0]

            # plot related
            self.plot = plot
            self.frame_files = []
            self.points = {}
            if self.plot:
                # initialize the route
                for i in range(self.num_robots):
                    self.points['x'+str(i+1)] = [self.start_robot_positions[i][0]]
                    self.points['y'+str(i+1)] = [self.start_robot_positions[i][1]]

        else:
            self.random_seed = random_seed
            self.num_robots = num_robots  
            if self.test:
                self.adj_dir = test_adj_path
            else:
                self.adj_dir = train_adj_path
            self.adj_list = os.listdir(self.adj_dir)
            self.adj_list.sort(reverse=False)
            self.adj_index = map_index % np.size(self.adj_list)
            self.graph, self.node_num, self.adjacent_matrix, self.network_adjacent_matrix, self.node_feature, self.next_node, self.start_robot_indexes, self.start_target_index, self.exit_indexes = self.import_adj_matrix(self.adj_dir + '/' + self.adj_list[self.adj_index], self.adj_index)
            self.robot_indexes = []
            for k in range(num_robots):
                self.robot_indexes.append(self.start_robot_indexes[k])
            self.target_index = self.start_target_index

            self.stepi = 0

            # if self.test:
            #     self.reference_policy = np.load(str(test_dp_path) + '/pursuer_policy_'+str(self.adj_index)+'.npy')
            # else:
            #     self.reference_policy = np.load(str(train_dp_path) + '/pursuer_policy_'+str(self.adj_index)+'.npy')
            
            # if FIXED_OPPONENT == False:
            #     if self.test:
            #         self.opponent_policy = np.load(str(test_dp_path) + '/opponent_policy_'+str(self.adj_index)+'.npy')
            #     else:
            #         self.opponent_policy = np.load(str(train_dp_path) + '/opponent_policy_'+str(self.adj_index)+'.npy')

    def find_index_from_coords(self, position):
        index = np.argmin(np.linalg.norm(self.node_coords - position, axis=1))
        assert self.node_num == np.linalg.norm(self.node_coords - position, axis=1).shape[0]
        return index
    
    def begin(self):
        # file_path = str(self.map_list[self.map_index])

        if EXIT_NUM > 0:
            self.start_robot_positions, self.start_robot_indexes, self.start_target_position, self.start_target_index, self.node_coords, self.graph, self.node_feature, self.start_robot_indexes, self.start_target_index, self.adjacent_matrix, self.network_adjacent_matrix, self.next_node, self.exit_indexes = self.graph_generator.generate_graph(
                self.ground_truth, self.random_seed)
            
        else:
            self.robot_positions, self.robot_indexes, self.target_position, self.target_index, self.node_coords, self.graph, self.node_feature, self.robot_indexes, self.target_index, self.adjacent_matrix, self.network_adjacent_matrix = self.graph_generator.generate_graph(
                self.ground_truth, self.random_seed)
        self.edge_mask = torch.from_numpy(self.adjacent_matrix).float().unsqueeze(0)
        self.node_num = self.node_coords.shape[0]

    # def begin(self):
    #     file_path = str(self.map_list[self.map_index])

    #     self.node_coords, self.graph, self.node_feature, self.robot_indexes, self.target_index, self.adjacent_matrix, self.network_adjacent_matrix = self.graph_generator.generate_graph(
    #         self.robot_positions, self.target_position, self.ground_truth, file_path, self.map_index)
        
    #     self.node_num = self.node_coords.shape[0]

    #     if self.test:
    #         if self.map_index % 2 == 0:
    #             self.reference_policy = np.load(str(test_small_map_path) + '/pursuer_policy_'+str(self.map_index)+'.npy')
    #         else:
    #             self.reference_policy = np.load(str(test_big_map_path) + '/pursuer_policy_'+str(self.map_index)+'.npy')

    #         # self.reference_policy = np.load(str(test_map_path) + '/pursuer_policy_'+str(self.map_index)+'.npy')
    #     else:
    #         if self.map_index % 2 == 0:
    #             self.reference_policy = np.load(str(train_small_map_path) + '/pursuer_policy_'+str(self.map_index)+'.npy')
    #         else:
    #             self.reference_policy = np.load(str(train_big_map_path) + '/pursuer_policy_'+str(self.map_index)+'.npy')
      
    #         # self.reference_policy = np.load(str(train_map_path) + '/pursuer_policy_'+str(self.map_index)+'.npy')
        
    #     if FIXED_OPPONENT == False:
    #         # self.opponent_policy = self.graph_generator.generate_opponent_policy(self.map_index)
    #         # self.opponent_policy = Global_config.opponent_policy[self.map_index]
    #         if self.test:
    #             if self.map_index % 2 == 0:
    #                 self.opponent_policy = np.load(str(test_small_map_path) + '/opponent_policy_'+str(self.map_index)+'.npy')
    #             else:
    #                 self.opponent_policy = np.load(str(test_big_map_path) + '/opponent_policy_'+str(self.map_index)+'.npy')

    #             # self.opponent_policy = np.load(str(test_map_path) + '/opponent_policy_'+str(self.map_index)+'.npy')
    #         else:
    #             if self.map_index % 2 == 0:
    #                 self.opponent_policy = np.load(str(train_small_map_path) + '/opponent_policy_'+str(self.map_index)+'.npy')
    #             else:
    #                 self.opponent_policy = np.load(str(train_big_map_path) + '/opponent_policy_'+str(self.map_index)+'.npy')

    #             # self.opponent_policy = np.load(str(train_map_path) + '/opponent_policy_'+str(self.map_index)+'.npy')

    def resetenv(self):
        self.robot_indexes = copy.deepcopy(self.start_robot_indexes)
        self.target_index = copy.deepcopy(self.start_target_index)
        self.stepi = 0
        if self.input_type == 'map':
            self.robot_positions = copy.deepcopy(self.start_robot_positions)
            self.target_position = copy.deepcopy(self.start_target_position)
            self.frame_files = []
            self.node_feature = self.graph_generator.update_graph(self.robot_indexes, self.target_index)
            self.points = {}
            if self.plot:
                # initialize the route
                for i in range(self.num_robots):
                    self.points['x'+str(i+1)] = [self.start_robot_positions[i][0]]
                    self.points['y'+str(i+1)] = [self.start_robot_positions[i][1]]
        # self.node_coords, self.graph, self.node_feature, self.adjacent_matrix, self.network_adjacent_matrix = None, None, None, None, None
        
        # self.node_num = 0
        # self.begin()
    
            
        else:
            node_feature = []
            for index in range(self.node_num):
                feature = []
                feature.append(self.network_adjacent_matrix[index][self.target_index])
                for robot in range(self.num_robots):
                    feature.append(self.network_adjacent_matrix[index][robot])
                if EXIT_NUM > 0:
                    for exit in self.exit_indexes:
                        feature.append(self.network_adjacent_matrix[index][exit])
                # print(index, feature)
                node_feature.append(feature)
            node_feature = np.array(node_feature)
       

    # 待修改
    def step(self, next_position, robot_index):
        if self.input_type == 'map':
            if robot_index <= N_ROBOTS:
                next_node_index = self.find_index_from_coords(next_position)
                self.robot_positions[robot_index] = next_position
                self.robot_indexes[robot_index] = next_node_index

                if self.plot:
                    self.points['x'+str(robot_index+1)].append(self.robot_positions[robot_index][0])
                    self.points['y'+str(robot_index+1)].append(self.robot_positions[robot_index][1])
            
                # check if done
                done, escaped = self.check_done(self.robot_indexes)
                reward = 0
                if self.test:
                    if escaped:
                        reward = -1
                    if done:
                        reward = 1
                    # else:
                    #     reward = 0

                else:
                    if escaped:
                        reward = -1
                    if done:
                        reward = 0

                # update the graph
                self.node_feature = self.graph_generator.update_graph(self.robot_indexes, self.target_index)

                return reward, done, self.robot_positions
            else:
                if FIXED_OPPONENT == False:
                    self.target_index = next_position
                    self.target_position = self.node_coords[self.target_index]
                self.stepi += 1

                # update the graph
                self.node_feature = self.graph_generator.update_graph(self.robot_indexes, self.target_index)
        else:
            if robot_index <= N_ROBOTS:
                next_index = next_position
                self.robot_indexes[robot_index] = int(next_index)
                node_feature = []

                for index in range(self.node_num):
                    feature = []
                    feature.append(self.network_adjacent_matrix[index][self.target_index])
                    for robot in range(self.num_robots):
                        feature.append(self.network_adjacent_matrix[index][self.robot_indexes[robot]])
                    if EXIT_NUM > 0:
                        for exit in self.exit_indexes:
                            feature.append(self.network_adjacent_matrix[index][exit])
                    # print(index, feature)
                    node_feature.append(feature)
                self.node_feature = np.array(node_feature)
                
            else:
                self.target_index = next_position
                self.stepi += 1
                node_feature = []
                for index in range(self.node_num):
                    feature = []
                    feature.append(self.network_adjacent_matrix[index][self.target_index])
                    for robot in range(self.num_robots):
                        feature.append(self.network_adjacent_matrix[index][self.robot_indexes[robot]])
                    if EXIT_NUM > 0:
                        for exit in self.exit_indexes:
                            feature.append(self.network_adjacent_matrix[index][exit])
                    # print(index, feature)
                    node_feature.append(feature)
                self.node_feature = np.array(node_feature)

    def import_ground_truth(self, map_index):
        # occupied 1, free 255, unexplored 127
        ground_truth = (io.imread(map_index, 1) * 255).astype(int)
        ground_truth = (ground_truth > 150)
        ground_truth = ground_truth * 254 + 1

        return ground_truth
    
    def import_adj_matrix(self, file_path, index):
        with open(file_path, 'r') as f:
            node_num = int(f.readline().strip())
            adj_matrix = []
            for line in f:
                adj_matrix.append(list(map(int, line.split())))
            adj_matrix = np.array(adj_matrix)
        graph = nx.from_numpy_array(adj_matrix)
        # print('edge num: ', graph.number_of_edges(), file_path)
        # nx.draw(graph, with_labels=True, node_color='red', node_size=100)
        # plt.savefig(f'fig/{file_path}.png')
        # network_adjacent_matrix = nx.floyd_warshall_numpy(graph)
        

        # next_node = [[j for j in range(node_num)] for i in range(node_num)]
        # for i in range(node_num):
        #     for j in range(node_num):
        #         for k in range(node_num):
        #             if network_adjacent_matrix[i][k] == 1 and network_adjacent_matrix[i][j] == network_adjacent_matrix[k][j] + 1:
        #                 next_node[i][j] = k
        # next_node = np.array(next_node)

        # if self.random_seed != None: 
        #     random.seed(self.random_seed)
        #     robot_index = random.randint(0, node_num-1)
        #     random.seed(self.random_seed + 10)
        #     target_index = random.randint(0, node_num-1)
        # else:
        #     # count = 0
        #     # while True:
        #     #     count += 1
        #     #     robot_index = random.randint(0, node_num-1)
        #     #     target_index = random.randint(0, node_num-1)
        #     #     dist = nx.shortest_path_length(graph, source=robot_index, target=target_index)
        #     #     # print('dist: ', dist)
        #     #     if dist > 5 or count > 1000:
        #     #         break
        #     if EXIT_NUM > 0: 
        #         cnt = 0        
        #         while True:
        #             cnt += 1
        #             print('-------------', cnt)
        #             # 随机选择出口
        #             random.seed(42)
        #             exit_indexes = random.sample(range(node_num), EXIT_NUM)

        #             random.seed()
        #             # 选择追捕者位置，排除出口节点
        #             candidates = [index for index in range(node_num) if index not in exit_indexes]
        #             robot_indexes = random.sample(candidates, N_ROBOTS)

        #             # 固定逃逸者到最近出口的距离为指定值
        #             min_dist = MIN_EVADOR_EXIT_DIST  # 设置逃逸者到最近出口的固定距离

        #             # 筛选符合条件的逃逸者节点
        #             valid_target_candidates = [
        #                 node for node in candidates
        #                 if any(
        #                     network_adjacent_matrix[node][exit_index] == min_dist
        #                     for exit_index in exit_indexes
        #                 )
        #             ]

        #             if not valid_target_candidates:
        #                 print("No valid target nodes found for the desired distance. Adjust parameters.")
        #                 continue  # 跳过当前循环，重新选择出口和追捕者

        #             # 随机选择一个目标节点
        #             target_index = random.choice(valid_target_candidates)

        #             # 找出逃逸者最近的出口
        #             closest_exit_index = min(
        #                 exit_indexes,
        #                 key=lambda exit_index: network_adjacent_matrix[target_index][exit_index]
        #             )
        #             closest_exit_distance = network_adjacent_matrix[target_index][closest_exit_index]

        #             # 检查条件：最近出口距离必须等于 DESIRED_DISTANCE
        #             if closest_exit_distance != min_dist:
        #                 continue

        #             # # 检查条件：最近出口是否有追捕者比逃逸者更接近
        #             # closer_robots = [
        #             #     robot_index for robot_index in robot_indexes
        #             #     if self.network_adjacent_matrix[robot_index][closest_exit_index] < closest_exit_distance
        #             # ]

        #             # if not closer_robots:  # 如果没有追捕者比逃逸者更接近最近出口
        #             #     continue


        #             # 检查条件 2：每个距离逃逸者小于 10 的出口，都有追捕者更近
        #             valid_exits = [
        #                 exit_index for exit_index in exit_indexes
        #                 if network_adjacent_matrix[target_index][exit_index] <= 10
        #             ]

        #             all_exits_satisfied = True
        #             for exit_index in valid_exits:
        #                 closer_robots = [
        #                     robot_index for robot_index in robot_indexes
        #                     if network_adjacent_matrix[robot_index][exit_index] < network_adjacent_matrix[target_index][exit_index]
        #                 ]
        #                 if not closer_robots:
        #                     all_exits_satisfied = False
        #                     break

        #             if not all_exits_satisfied:
        #                 continue

        #             # 如果满足所有条件，退出循环
        #             break
                    
        #     else:
        #         raise ValueError("Must Have Exit")
        # print('file', file_path, 'exit: ', exit_indexes, 'pursuer: ', robot_indexes, 'evader: ', target_index)
        
        if self.test == True:
            data = np.load('init/test.npz')
            network_adjacent_matrix = np.load('init/test/'+str(index)+'.npz')['dist']
            next_node = np.load('init/test/'+str(index)+'.npz')['next']
        else:
            data = np.load('init/train.npz')
            network_adjacent_matrix = np.load('init/train/'+str(index)+'.npz')['dist']
            next_node = np.load('init/train/'+str(index)+'.npz')['next']
        assert network_adjacent_matrix.shape[0] == adj_matrix.shape[0]
        
        exit_indexes = list(data['exit'][index])
        robot_indexes = list(data['robot'][index])
        target_index = data['target'][index]
        # print('file', file_path, 'index: ', index, 'exit: ', exit_indexes, 'pursuer: ', robot_indexes, 'evader: ', target_index)

        adj_matrix = 1 - adj_matrix
        
        node_feature = []
        for index in range(node_num):
            feature = []
            feature.append(network_adjacent_matrix[index][target_index])
            for robot in range(self.num_robots):
                feature.append(network_adjacent_matrix[index][robot])
            if EXIT_NUM > 0:
                for exit in exit_indexes:
                    feature.append(network_adjacent_matrix[index][exit])
            # print(index, feature)
            node_feature.append(feature)
        node_feature = np.array(node_feature)

        return graph, node_num, adj_matrix, network_adjacent_matrix, node_feature, next_node, robot_indexes, target_index, exit_indexes
    
    def free_cells(self):
        index = np.where(self.ground_truth == 255)
        free = np.asarray([index[1], index[0]]).T
        return free

    def check_done(self, node_indexes):
        # dist_threshold = self.graph_generator.dist_threshold
        cnt = 0
        for k in range(self.num_robots):
            # dist = np.linalg.norm(self.robot_positions[k] - self.target_position)
            if node_indexes[k] == self.target_index:
                cnt += 1
        done = False
        if cnt >= 1:
            done = True

        escaped = False
        if EXIT_NUM > 0:
            if self.target_index in self.exit_indexes:
                escaped = True
        
        # return done
        return done, escaped

    # def calculate_reward(self, dist):
    #     reward = 0
    #     return reward
    
    def evaluate_exploration_rate(self):
        rate = np.sum(self.combined_belief == 255) / np.sum(self.ground_truth == 255)
        return rate

    def calculate_new_free_area(self, old_robot_belief, robot_belief):
        old_free_area = old_robot_belief == 255
        current_free_area = robot_belief == 255
        new_free_area = (current_free_area.astype(int) - old_free_area.astype(int)) * 255
        new_free_area = np.where(new_free_area == -255, 0, new_free_area)
        return new_free_area

    def calculate_path_length(self, path):
        dist = 0
        start = path[0]
        end = path[-1]
        for index in path:
            if index == end:
                break
            dist += np.linalg.norm(self.node_coords[start] - self.node_coords[index])
            start = index
        return dist

    def plot_env(self, n, path, step):
        plt.switch_backend('agg')
        colors = ['r', 'g', 'y', 'm', 'skyblue']
        plt.cla()
        plt.suptitle('')
        plt.imshow(self.ground_truth, cmap='gray')
        plt.axis('off')
        # plt.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0))
        for i in range(len(self.graph_generator.x)):
            plt.plot(self.graph_generator.x[i], self.graph_generator.y[i], 'orange', zorder=1)  # plot edges will take long time
        
        for i in range(self.num_robots):    
            plt.scatter(self.points['x'+str(i+1)][-1], self.points['y'+str(i+1)][-1], c=colors[i], s=200-i*20, zorder=7)

        if EXIT_NUM > 0:
            # if self.state == 2:
            #     state = 'Escaped'
            # elif self.state == 1:
            #     state = 'Catched'
            # else:
            #     state = 'Not escaped or catched'

            plt.suptitle('Total step: {}'.format(self.stepi))
            # plt.savefig('{}/{}_{}_{}_samples.png'.format(path, n, step, self.state, dpi=300))
            frame = '{}/{}_{}_samples.png'.format(path, n, step)
            for exit_index in self.exit_indexes:
                exit_position = self.node_coords[exit_index]
                plt.scatter(exit_position[0], exit_position[1], c='lime', marker='*', s=250, zorder=6)
    
        else:
            plt.suptitle('Total step: {}'.format(self.stepi))
            # plt.savefig('{}/{}_{}_samples.png'.format(path, n, step, dpi=300))
            frame = '{}/{}_{}_samples.png'.format(path, n, step)
        
        plt.scatter(self.node_coords[:, 0], self.node_coords[:, 1], c='darkblue', zorder=5)
        plt.scatter(self.node_coords[self.target_index, 0], self.node_coords[self.target_index, 1], s=100, marker='s', c='c', zorder=10)
        plt.tight_layout()
        plt.savefig(frame)
        # plt.show()
        self.frame_files.append(frame)
        plt.close()

    def get_opponent_position(self, robot_index, opponent_index):
        # print("opponent index:", opponent_index, "robot index:", robot_index, 'map index: ', self.map_index, 'map file: ', self.map_dir + '/' + self.map_list[self.map_index])
        # next_opponent_position = self.opponent_policy[int(opponent_index)][robot_index[0]][robot_index[1]][robot_index[2]]
        next_opponent_position = self.opponent_policy[int(opponent_index)][robot_index[0]][robot_index[1]]
       
        return next_opponent_position